import timm
from typing import Callable, Any, Optional, Tuple, List
from . import utils


# NOTICE: MUST list the modules from inner to outer
Inc_Res_V2_Default_DHF_Modules = [
    "repeat.6.branch2.0.conv",
    "repeat.6.branch2.1.conv",
    "repeat.6.branch2.2.conv",
    "repeat.6.conv2d",
    "repeat.7.branch2.0.conv",
    "repeat.7.branch2.1.conv",
    "repeat.7.branch2.2.conv",
    "repeat.7.conv2d",
    "repeat.8.branch2.0.conv",
    "repeat.8.branch2.1.conv",
    "repeat.8.branch2.2.conv",
    "repeat.8.conv2d",
    "repeat.9.branch2.0.conv",
    "repeat.9.branch2.1.conv",
    "repeat.9.branch2.2.conv",
    "repeat.9.conv2d",
    "mixed_6a.branch1.0.conv",
    "mixed_6a.branch1.1.conv",
    "mixed_6a.branch1.2.conv",
    "repeat_1.0.branch1.0.conv",
    "repeat_1.0.branch1.1.conv",
    "repeat_1.0.branch1.2.conv",
    "repeat_1.0.conv2d",
    "repeat_1.1.branch1.0.conv",
    "repeat_1.1.branch1.1.conv",
    "repeat_1.1.branch1.2.conv",
    "repeat_1.1.conv2d",
    "repeat_1.2.branch1.0.conv",
    "repeat_1.2.branch1.1.conv",
    "repeat_1.2.branch1.2.conv",
    "repeat_1.2.conv2d",
    "repeat_1.3.branch1.0.conv",
    "repeat_1.3.branch1.1.conv",
    "repeat_1.3.branch1.2.conv",
    "repeat_1.3.conv2d",
    "repeat_1.4.branch1.0.conv",
    "repeat_1.4.branch1.1.conv",
    "repeat_1.4.branch1.2.conv",
    "repeat_1.4.conv2d",
    "repeat_1.5.branch1.0.conv",
    "repeat_1.5.branch1.1.conv",
    "repeat_1.5.branch1.2.conv",
    "repeat_1.5.conv2d",
    "repeat_1.6.branch1.0.conv",
    "repeat_1.6.branch1.1.conv",
    "repeat_1.6.branch1.2.conv",
    "repeat_1.6.conv2d",
    "repeat_1.7.branch1.0.conv",
    "repeat_1.7.branch1.1.conv",
    "repeat_1.7.branch1.2.conv",
    "repeat_1.7.conv2d",
    "repeat_1.8.branch1.0.conv",
    "repeat_1.8.branch1.1.conv",
    "repeat_1.8.branch1.2.conv",
    "repeat_1.8.conv2d",
    "repeat_1.9.branch1.0.conv",
    "repeat_1.9.branch1.1.conv",
    "repeat_1.9.branch1.2.conv",
    "repeat_1.9.conv2d",
    "repeat_1.10.branch1.0.conv",
    "repeat_1.10.branch1.1.conv",
    "repeat_1.10.branch1.2.conv",
    "repeat_1.10.conv2d",
    "repeat_1.11.branch1.0.conv",
    "repeat_1.11.branch1.1.conv",
    "repeat_1.11.branch1.2.conv",
    "repeat_1.11.conv2d",
    "repeat_1.12.branch1.0.conv",
    "repeat_1.12.branch1.1.conv",
    "repeat_1.12.branch1.2.conv",
    "repeat_1.12.conv2d",
    "repeat_1.13.branch1.0.conv",
    "repeat_1.13.branch1.1.conv",
    "repeat_1.13.branch1.2.conv",
    "repeat_1.13.conv2d",
    "repeat_1.14.branch1.0.conv",
    "repeat_1.14.branch1.1.conv",
    "repeat_1.14.branch1.2.conv",
    "repeat_1.14.conv2d",
    "repeat_1.15.branch1.0.conv",
    "repeat_1.15.branch1.1.conv",
    "repeat_1.15.branch1.2.conv",
    "repeat_1.15.conv2d",
    "repeat_1.16.branch1.0.conv",
    "repeat_1.16.branch1.1.conv",
    "repeat_1.16.branch1.2.conv",
    "repeat_1.16.conv2d",
    "repeat_1.17.branch1.0.conv",
    "repeat_1.17.branch1.1.conv",
    "repeat_1.17.branch1.2.conv",
    "repeat_1.17.conv2d",
    "repeat_1.18.branch1.0.conv",
    "repeat_1.18.branch1.1.conv",
    "repeat_1.18.branch1.2.conv",
    "repeat_1.18.conv2d",
    "repeat_1.19.branch1.0.conv",
    "repeat_1.19.branch1.1.conv",
    "repeat_1.19.branch1.2.conv",
    "repeat_1.19.conv2d",
    "mixed_7a.branch2.0.conv",
    "mixed_7a.branch2.1.conv",
    "mixed_7a.branch2.2.conv",
    "repeat_2.0.branch1.0.conv",
    "repeat_2.0.branch1.1.conv",
    "repeat_2.0.branch1.2.conv",
    "repeat_2.0.conv2d",
    "repeat_2.1.branch1.0.conv",
    "repeat_2.1.branch1.1.conv",
    "repeat_2.1.branch1.2.conv",
    "repeat_2.1.conv2d",
    "repeat_2.2.branch1.0.conv",
    "repeat_2.2.branch1.1.conv",
    "repeat_2.2.branch1.2.conv",
    "repeat_2.2.conv2d",
    "repeat_2.3.branch1.0.conv",
    "repeat_2.3.branch1.1.conv",
    "repeat_2.3.branch1.2.conv",
    "repeat_2.3.conv2d",
    "repeat_2.4.branch1.0.conv",
    "repeat_2.4.branch1.1.conv",
    "repeat_2.4.branch1.2.conv",
    "repeat_2.4.conv2d",
    "repeat_2.5.branch1.0.conv",
    "repeat_2.5.branch1.1.conv",
    "repeat_2.5.branch1.2.conv",
    "repeat_2.5.conv2d",
    "repeat_2.6.branch1.0.conv",
    "repeat_2.6.branch1.1.conv",
    "repeat_2.6.branch1.2.conv",
    "repeat_2.6.conv2d",
    "repeat_2.7.branch1.0.conv",
    "repeat_2.7.branch1.1.conv",
    "repeat_2.7.branch1.2.conv",
    "repeat_2.7.conv2d",
    "repeat_2.8.branch1.0.conv",
    "repeat_2.8.branch1.1.conv",
    "repeat_2.8.branch1.2.conv",
    "repeat_2.8.conv2d",
    "block8.branch1.0.conv",
    "block8.branch1.1.conv",
    "block8.branch1.2.conv",
    "block8.conv2d",
]


def inc_res_v2(*args, **kwargs):
    model = timm.create_model('inception_resnet_v2', pretrained=True)
    return model


def dhf_inc_res_v2(mixup_weight_max: float, random_keep_prob: float, dhf_modules = None, weights = None, progress: bool = True, **kwargs: Any):
    dhf_model = timm.create_model('inception_resnet_v2', pretrained=True)
    if dhf_modules is None:
        dhf_modules = Inc_Res_V2_Default_DHF_Modules
    utils.convert_to_DHF_model_inplace_(dhf_model, mixup_weight_max=mixup_weight_max, random_keep_prob=random_keep_prob, dhf_modules=dhf_modules)
    return dhf_model